{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7dd76610",
   "metadata": {},
   "source": [
    "# Python SDK for Lineage\n",
    "\n",
    "Full docs here https://docs.open-metadata.org/sdk/python/ingestion/lineage"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb409c67",
   "metadata": {},
   "source": [
    "## 1. Preparing the Client\n",
    "\n",
    "Here we will define how to connect to the OpenMetadata API, using the JWT token"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "60ac8e18",
   "metadata": {},
   "outputs": [],
   "source": [
    "from metadata.generated.schema.entity.services.connections.metadata.openMetadataConnection import (\n",
    "    OpenMetadataConnection,\n",
    ")\n",
    "from metadata.generated.schema.security.client.openMetadataJWTClientConfig import (\n",
    "    OpenMetadataJWTClientConfig,\n",
    ")\n",
    "from metadata.ingestion.ometa.ometa_api import OpenMetadata\n",
    "\n",
    "server_config = OpenMetadataConnection(\n",
    "    hostPort=\"http://localhost:8585/api\",\n",
    "    authProvider=\"openmetadata\",\n",
    "    securityConfig=OpenMetadataJWTClientConfig(\n",
    "        jwtToken=\"eyJraWQiOiJHYjM4OWEtOWY3Ni1nZGpzLWE5MmotMDI0MmJrOTQzNTYiLCJ0eXAiOiJKV1QiLCJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJhZG1pbiIsImlzQm90IjpmYWxzZSwiaXNzIjoib3Blbi1tZXRhZGF0YS5vcmciLCJpYXQiOjE2NjM5Mzg0NjIsImVtYWlsIjoiYWRtaW5Ab3Blbm1ldGFkYXRhLm9yZyJ9.tS8um_5DKu7HgzGBzS1VTA5uUjKWOCU0B_j08WXBiEC0mr0zNREkqVfwFDD-d24HlNEbrqioLsBuFRiwIWKc1m_ZlVQbG7P36RUxhuv2vbSp80FKyNM-Tj93FDzq91jsyNmsQhyNv_fNr3TXfzzSPjHt8Go0FMMP66weoKMgW2PbXlhVKwEuXUHyakLLzewm9UMeQaEiRzhiTMU3UkLXcKbYEJJvfNFcLwSl9W8JCO_l0Yj3ud-qt_nQYEZwqW6u5nfdQllN133iikV4fM5QZsMCnm8Rq1mvLR0y9bmJiD7fwM1tmJ791TUWqmKaTnP49U493VanKpUAfzIiOiIbhg\"\n",
    "    ),\n",
    ")\n",
    "metadata = OpenMetadata(server_config)\n",
    "\n",
    "assert metadata.health_check()  # Will fail if we cannot reach the server"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2afb1023",
   "metadata": {},
   "source": [
    "## 2. Getting the tables\n",
    "\n",
    "You can do any operation using the OpenMetadata client or directly using the API. Now, we'll query the two tables we want to use to add lineage between them.\n",
    "\n",
    "We want to link the `actor` and `film_actor` tables, knowing that the ID is a relationship between both Entities:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9c6fe791",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "metadata.generated.schema.entity.data.table.Table"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from metadata.generated.schema.entity.data.table import Table\n",
    "\n",
    "actor_entity = metadata.get_by_name(entity=Table, fqn=\"demo_pg.postgres.public.actor\")\n",
    "film_actor_entity = metadata.get_by_name(entity=Table, fqn=\"demo_pg.postgres.public.film_actor\")\n",
    "\n",
    "type(actor_entity)  # Everything is typed :)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c472c90",
   "metadata": {},
   "source": [
    "## 3. Add Table Lineage\n",
    "\n",
    "The first step will be to create a lineage relationship between both tables:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "97808594",
   "metadata": {},
   "outputs": [],
   "source": [
    "from metadata.generated.schema.type.entityReference import EntityReference\n",
    "\n",
    "from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest\n",
    "from metadata.generated.schema.type.entityLineage import EntitiesEdge\n",
    "\n",
    "add_lineage_request = AddLineageRequest(\n",
    "    description=\"test lineage\",\n",
    "    edge=EntitiesEdge(\n",
    "        fromEntity=EntityReference(id=actor_entity.id, type=\"table\"),\n",
    "        toEntity=EntityReference(id=film_actor_entity.id, type=\"table\"),\n",
    "    ),\n",
    ")\n",
    "\n",
    "metadata.add_lineage(data=add_lineage_request)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1f50d46",
   "metadata": {},
   "source": [
    "## 4. Fetching lineage\n",
    "\n",
    "We can also gather the results at any time with:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "70eb0b3d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'entity': {'id': '1eb9b8e6-b538-4a8f-86ae-d49bd06a56b6',\n",
       "  'type': 'table',\n",
       "  'name': 'film_actor',\n",
       "  'fullyQualifiedName': 'demo_pg.postgres.public.film_actor',\n",
       "  'deleted': False,\n",
       "  'href': 'http://localhost:8585/api/v1/tables/1eb9b8e6-b538-4a8f-86ae-d49bd06a56b6'},\n",
       " 'nodes': [{'id': 'b3479b8a-b8ff-4368-a368-9614799969cd',\n",
       "   'type': 'table',\n",
       "   'name': 'actor',\n",
       "   'fullyQualifiedName': 'demo_pg.postgres.public.actor',\n",
       "   'deleted': False,\n",
       "   'href': 'http://localhost:8585/api/v1/tables/b3479b8a-b8ff-4368-a368-9614799969cd'}],\n",
       " 'upstreamEdges': [{'fromEntity': 'b3479b8a-b8ff-4368-a368-9614799969cd',\n",
       "   'toEntity': '1eb9b8e6-b538-4a8f-86ae-d49bd06a56b6',\n",
       "   'lineageDetails': {'sqlQuery': 'SELECT * FROM AWESOME',\n",
       "    'columnsLineage': [{'fromColumns': ['demo_pg.postgres.public.actor.actor_id'],\n",
       "      'toColumn': 'demo_pg.postgres.public.film_actor.actor_id'}]}}],\n",
       " 'downstreamEdges': []}"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "metadata.get_lineage_by_name(\n",
    "    entity=Table,\n",
    "    fqn=\"demo_pg.postgres.public.film_actor\",\n",
    "    # Tune this to control how far in the lineage graph to go\n",
    "    up_depth=1,\n",
    "    down_depth=1\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13052037",
   "metadata": {},
   "source": [
    "## 5. Spice things up with Column Level Lineage\n",
    "\n",
    "Which is just an extra step when creating lineage."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "a6828f36",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'entity': {'id': 'b3479b8a-b8ff-4368-a368-9614799969cd',\n",
       "  'type': 'table',\n",
       "  'name': 'actor',\n",
       "  'fullyQualifiedName': 'demo_pg.postgres.public.actor',\n",
       "  'deleted': False,\n",
       "  'href': 'http://localhost:8585/api/v1/tables/b3479b8a-b8ff-4368-a368-9614799969cd'},\n",
       " 'nodes': [{'id': '1eb9b8e6-b538-4a8f-86ae-d49bd06a56b6',\n",
       "   'type': 'table',\n",
       "   'name': 'film_actor',\n",
       "   'fullyQualifiedName': 'demo_pg.postgres.public.film_actor',\n",
       "   'deleted': False,\n",
       "   'href': 'http://localhost:8585/api/v1/tables/1eb9b8e6-b538-4a8f-86ae-d49bd06a56b6'}],\n",
       " 'upstreamEdges': [],\n",
       " 'downstreamEdges': [{'fromEntity': 'b3479b8a-b8ff-4368-a368-9614799969cd',\n",
       "   'toEntity': '1eb9b8e6-b538-4a8f-86ae-d49bd06a56b6',\n",
       "   'lineageDetails': {'sqlQuery': 'SELECT * FROM AWESOME',\n",
       "    'columnsLineage': [{'fromColumns': ['demo_pg.postgres.public.actor.actor_id'],\n",
       "      'toColumn': 'demo_pg.postgres.public.film_actor.actor_id'}]}}]}"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from metadata.generated.schema.type.entityLineage import (\n",
    "    ColumnLineage,\n",
    "    EntitiesEdge,\n",
    "    LineageDetails,\n",
    ")\n",
    "\n",
    "column_lineage = ColumnLineage(\n",
    "    fromColumns=[\"demo_pg.postgres.public.actor.actor_id\"],\n",
    "    toColumn=\"demo_pg.postgres.public.film_actor.actor_id\"\n",
    ")\n",
    "\n",
    "lineage_details = LineageDetails(\n",
    "    sqlQuery=\"SELECT * FROM AWESOME\",\n",
    "    columnsLineage=[column_lineage],\n",
    ")\n",
    "\n",
    "add_lineage_request = AddLineageRequest(\n",
    "    edge=EntitiesEdge(\n",
    "        fromEntity=EntityReference(id=actor_entity.id, type=\"table\"),\n",
    "        toEntity=EntityReference(id=film_actor_entity.id, type=\"table\"),\n",
    "        lineageDetails=lineage_details,\n",
    "    ),\n",
    ")\n",
    "\n",
    "metadata.add_lineage(data=add_lineage_request)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b53733f2",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
